"""
ConnLLM OpenRouter思考过程能力V2实现
提供从OpenRouter API响应中提取思考过程的功能
"""

from typing import Dict, Any, List, Optional, Union, Pattern
import re
import logging
from abc import ABC, abstractmethod

from ...core_v2.capability import ModelCapability

logger = logging.getLogger("connllm.implementations.openrouter.thinking_v2")

class ThinkingExtractor(ABC):
	"""思考过程提取器基类"""
	
	@abstractmethod
	def extract(self, text: str) -> Optional[str]:
		"""
		从文本中提取思考过程
		
		Args:
			text: 响应文本
			
		Returns:
			提取的思考过程，如果没有则返回None
		"""
		pass
	
	@abstractmethod
	def clean(self, text: str) -> str:
		"""
		从文本中移除思考过程
		
		Args:
			text: 响应文本
			
		Returns:
			清理后的文本
		"""
		pass


class RegexThinkingExtractor(ThinkingExtractor):
	"""基于正则表达式的思考过程提取器"""
	
	def __init__(self, pattern: Union[str, Pattern], group_name: Optional[str] = None):
		"""
		初始化提取器
		
		Args:
			pattern: 正则表达式或已编译的正则表达式
			group_name: 捕获组名称，若为None则使用第一个捕获组
		"""
		if isinstance(pattern, str):
			self.pattern = re.compile(pattern, re.DOTALL)
		else:
			self.pattern = pattern
			
		self.group_name = group_name
	
	def extract(self, text: str) -> Optional[str]:
		"""
		从文本中提取思考过程
		
		Args:
			text: 响应文本
			
		Returns:
			提取的思考过程，如果没有则返回None
		"""
		if not text:
			return None
			
		match = self.pattern.search(text)
		if not match:
			return None
			
		if self.group_name:
			return match.group(self.group_name).strip()
		else:
			return match.group(1).strip()
	
	def clean(self, text: str) -> str:
		"""
		从文本中移除思考过程
		
		Args:
			text: 响应文本
			
		Returns:
			清理后的文本
		"""
		if not text:
			return ""
			
		return self.pattern.sub("", text).strip()


class ClaudeThinkingExtractor(RegexThinkingExtractor):
	"""Claude模型思考过程提取器"""
	
	def __init__(self):
		"""初始化提取器"""
		pattern = r"<thinking>(.*?)</thinking>"
		super().__init__(pattern)


class OpenAIThinkingExtractor(RegexThinkingExtractor):
	"""OpenAI模型思考过程提取器"""
	
	def __init__(self):
		"""初始化提取器"""
		# OpenAI可能使用多种格式，这里列出几种可能的格式
		patterns = [
			r"<internal>(.*?)</internal>",
			r"\[\[([^]]+)\]\]",
			r"<!--(.*?)-->",
			r"/\*(.*?)\*/",
			r"<\!--(.*?)--\>"  # 注意这个模式与前面略有不同，应对HTML注释的不同写法
		]
		
		# 创建组合正则模式
		combined = "(" + "|".join(patterns) + ")"
		self.pattern = re.compile(combined, re.DOTALL)
		
		# 跟踪单个模式，用于clean方法
		self.individual_patterns = [re.compile(p, re.DOTALL) for p in patterns]
	
	def extract(self, text: str) -> Optional[str]:
		"""
		从文本中提取思考过程
		
		Args:
			text: 响应文本
			
		Returns:
			提取的思考过程，如果没有则返回None
		"""
		if not text:
			return None
			
		# 尝试所有模式
		for i, pattern in enumerate(self.individual_patterns):
			match = pattern.search(text)
			if match:
				return match.group(1).strip()
				
		return None
	
	def clean(self, text: str) -> str:
		"""
		从文本中移除思考过程
		
		Args:
			text: 响应文本
			
		Returns:
			清理后的文本
		"""
		if not text:
			return ""
			
		result = text
		# 应用所有清理模式
		for pattern in self.individual_patterns:
			result = pattern.sub("", result)
			
		return result.strip()


class GenericThinkingExtractor(ThinkingExtractor):
	"""通用思考过程提取器，组合多种提取器"""
	
	def __init__(self):
		"""初始化提取器"""
		self.extractors = [
			ClaudeThinkingExtractor(),
			OpenAIThinkingExtractor(),
			# 可以添加更多提取器
		]
	
	def extract(self, text: str) -> Optional[str]:
		"""
		从文本中提取思考过程
		
		Args:
			text: 响应文本
			
		Returns:
			提取的思考过程，如果没有则返回None
		"""
		if not text:
			return None
			
		# 尝试所有提取器
		for extractor in self.extractors:
			result = extractor.extract(text)
			if result:
				return result
				
		return None
	
	def clean(self, text: str) -> str:
		"""
		从文本中移除思考过程
		
		Args:
			text: 响应文本
			
		Returns:
			清理后的文本
		"""
		if not text:
			return ""
			
		result = text
		# 应用所有清理方法
		for extractor in self.extractors:
			result = extractor.clean(result)
			
		return result.strip()


class OpenRouterThinkingV2(ModelCapability):
	"""OpenRouter思考过程能力V2实现"""
	
	def __init__(self, provider):
		"""
		初始化能力
		
		Args:
			provider: 提供商实例
		"""
		super().__init__(provider)
		self._init_extractors()
	
	def _init_extractors(self):
		"""初始化提取器"""
		self.extractor_map = {
			"anthropic": ClaudeThinkingExtractor(),
			"claude": ClaudeThinkingExtractor(),
			"openai": OpenAIThinkingExtractor(),
			"gpt": OpenAIThinkingExtractor(),
			"default": GenericThinkingExtractor()
		}
	
	def _get_extractor_for_model(self, model: str) -> ThinkingExtractor:
		"""
		获取模型对应的提取器
		
		Args:
			model: 模型名称
			
		Returns:
			提取器实例
		"""
		# 尝试匹配模型前缀
		for prefix, extractor in self.extractor_map.items():
			if model.lower().startswith(prefix.lower()):
				return extractor
		
		# 默认使用通用提取器
		return self.extractor_map["default"]
	
	def extract_thinking(self, response: Dict[str, Any]) -> Optional[str]:
		"""
		从响应中提取思考过程
		
		Args:
			response: API响应
			
		Returns:
			提取的思考过程，如果没有则返回None
		"""
		try:
			# 首先检查是否有思考过程是通过OpenRouter特定功能提供的
			if "openrouter" in response and "thinking" in response["openrouter"]:
				return response["openrouter"]["thinking"]
			
			# 从响应中提取文本内容
			text = self._extract_text_from_response(response)
			if not text:
				return None
				
			# 获取模型名称
			model = self._get_model_from_response(response)
			
			# 获取适合模型的提取器
			extractor = self._get_extractor_for_model(model)
			
			# 提取思考过程
			return extractor.extract(text)
				
		except Exception as e:
			logger.error(f"提取思考过程失败: {str(e)}")
			return None
	
	def should_remove_thinking(self) -> bool:
		"""
		是否应该移除思考过程
		
		Returns:
			是否应该移除
		"""
		# 从配置中获取是否移除思考过程的设置
		return self.provider.config_manager.get("remove_thinking", True)
	
	def clean_response(self, response: Dict[str, Any]) -> Dict[str, Any]:
		"""
		从响应中移除思考过程
		
		Args:
			response: API响应
			
		Returns:
			清理后的响应
		"""
		try:
			# 如果不应该移除思考过程，直接返回
			if not self.should_remove_thinking():
				return response
				
			# 从响应中提取文本内容
			if "choices" in response and len(response["choices"]) > 0:
				choice = response["choices"][0]
				
				if "message" in choice and "content" in choice["message"]:
					text = choice["message"]["content"]
					
					# 获取模型名称
					model = self._get_model_from_response(response)
					
					# 获取适合模型的提取器
					extractor = self._get_extractor_for_model(model)
					
					# 清理文本
					cleaned_text = extractor.clean(text)
					
					# 更新响应
					choice["message"]["content"] = cleaned_text
			
			return response
				
		except Exception as e:
			logger.error(f"清理思考过程失败: {str(e)}")
			return response
	
	def _extract_text_from_response(self, response: Dict[str, Any]) -> Optional[str]:
		"""
		从响应中提取文本内容
		
		Args:
			response: API响应
			
		Returns:
			提取的文本，如果没有则返回None
		"""
		try:
			if "choices" in response and len(response["choices"]) > 0:
				choice = response["choices"][0]
				
				if "message" in choice and "content" in choice["message"]:
					return choice["message"]["content"]
				
				if "text" in choice:
					return choice["text"]
		except Exception as e:
			logger.error(f"从响应中提取文本失败: {str(e)}")
			
		return None
	
	def _get_model_from_response(self, response: Dict[str, Any]) -> str:
		"""
		从响应中获取模型名称
		
		Args:
			response: API响应
			
		Returns:
			模型名称
		"""
		try:
			# 首先检查响应中是否有模型信息
			if "model" in response:
				return response["model"]
				
			# 然后检查OpenRouter特定字段
			if "openrouter" in response and "model" in response["openrouter"]:
				return response["openrouter"]["model"]
		except Exception as e:
			logger.error(f"从响应中获取模型名称失败: {str(e)}")
			
		# 默认使用提供商当前配置的模型
		return self.provider.config_manager.get_model()
